import torch
from models.lprload import chars
from cv2 import imdecode, resize
from numpy import fromfile, uint8, newaxis, argmax, transpose


def prediction(_predict):  # lprnet的标签解码,65解码为无。
    predict = _predict.cpu().detach().numpy()[0, :, :]
    pre_index = [65]
    label = ""
    for i in range(predict.shape[1]):
        pre_number = argmax(predict[:, i], axis=0)
        pre_index.append(65 if pre_number == pre_index[-1] else pre_number)
    for index in pre_index:
        if index != 65:
            label += chars[index]
    return label


def detect_image(_image_path=''):
    if not _image_path.endswith((".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff")):
        return None  # 格式不正确图片
    try:
        image = imdecode(fromfile(_image_path, dtype=uint8), -1)
        image = resize(image, (94, 24))
    except AttributeError:  # 无字符图片
        return False
    try:
        transpose((image.astype('float32') - 127.5) / 128.0, (2, 0, 1))
    except ValueError:  # 无字符图片
        return False
    image = torch.Tensor(image[newaxis, :]).to(device)
    pre = model(image)
    label = prediction(pre)
    return label


device = torch.device('cpu')
model = torch.load('weights/best/oldlpr.pth', map_location=device)  # 预先加载，最好激活一次，调用更快
if __name__ == '__main__':
    # 1.单张预测
    print(detect_image(r"F:\ccpd4\lpr\blue\皖A0J872.jpg"))
    # 3.文件夹（多张预测），逐张预测输出
    # img_path = r"../ccpd4/lpr/train"
    # for path in os.listdir(img_path):
    #     print(detect_image(img_path + '/' + path))
